package fem2.material;

import inf.jlinalg.BLAM;
import inf.jlinalg.IMatrix;
import inf.jlinalg.IVector;

import java.util.HashMap;
import java.util.Map;

import math2.MathUtilities;
import math2.MutableDouble;
import fem2.Debugger;
import fem2.Element;
import fem2.MaterialModel;
import fem2.Observer;
import fem2.Variable;
import fem2.element.StructuralElement;
import fem2.enu.EchoLevelType;
import fem2.enu.VariableName;

/**
 * topology optimization material model based on homogenization method
 * 
 * @author hbui
 * 
 */
public class TopOptMaterialModel extends MaterialModel implements Observer {

	private class TopOptVariable extends Variable {

		public TopOptVariable(int buffSize, int vdim) {
			super(buffSize);
			registerScalarVariable(VariableName.DENSITY, true);
			registerScalarVariable(VariableName.PHYSICAL_DENSITY, true);
			registerScalarVariable(VariableName.DCDXE, true);
			registerScalarVariable(VariableName.DVDXE, true);
		}
	}

	private ElasticMaterialModel mm;
	private Map<Element, TopOptVariable> elementVariables = new HashMap<Element, TopOptVariable>();
	private double Emin, E0;
	private double penal;
	private double volfrac;
	private TopOptFilter filter;
	private double densityChange = Double.MAX_VALUE;

	public TopOptMaterialModel(ElasticMaterialModel mm, double E, double Emin, double penal,
			double volfrac, TopOptFilter filter) {
		super(mm);
		this.mm = mm;
		this.E0 = E;
		this.Emin = Emin;
		this.penal = penal;
		this.volfrac = volfrac;
		this.filter = filter;
	}

	/**
	 * get the nested material model associated with this material model
	 * 
	 * @return
	 */
	public MaterialModel getNestedMaterialModel() {
		return mm;
	}

	@Override
	public void Initialize(Element e) {
		mm.Initialize(e);
		/*
		 * initialize host material model
		 */
		TopOptVariable v = new TopOptVariable(1, 1);
		v.getScalarVariable(VariableName.DENSITY).setValue(volfrac);
		v.getScalarVariable(VariableName.PHYSICAL_DENSITY).setValue(volfrac);
		v.save();
		elementVariables.put(e, v);
	}

	@Override
	public void InitializeSolutionStep(Element e) {
		mm.InitializeSolutionStep(e);
		/*
		 * modify stress of element
		 */
		StructuralElement se = (StructuralElement) e;
		double xe = getPhysicalDensity(e).getValue();
		double Ee = computeModulus(xe);
		double coeff = Ee / E0;
		for (int i = 0; i < se.getNumIntegrationPoints(); i++) {
			IVector sigma = se.getStress(i);
			BLAM.scale(coeff, sigma);
		}
	}

	@Override
	public void InitalizeNonlinearIteration(Element e) {
		mm.InitalizeNonlinearIteration(e);
	}

	@Override
	public void FinalizeNonlinearIteration(Element e) {
		mm.FinalizeNonlinearIteration(e);
		StructuralElement se = (StructuralElement) e;
		/*
		 * modify stress of element
		 */
		double xe = getPhysicalDensity(e).getValue();
		double Ee = computeModulus(xe);
		double coeff = Ee / E0;
		for (int i = 0; i < se.getNumIntegrationPoints(); i++) {
			IVector sigma = se.getStress(i);
			BLAM.scale(coeff, sigma);
		}
		/*
		 * sensitivity analysis
		 */
		// double ce = 2 * se.computeStrainEnergy();
		double ce = 2 * mm.computeStrainEnergy(se);
		double dEe = computeDModulus(xe);
		double dc = -(dEe / Ee) * ce;
		double dv = e.getMeshPart().computeVolume();
		if (e.getEchoLevel().contain(EchoLevelType.DEBUG)) {
			Debugger.watch("ce = " + ce);
			Debugger.watch("dc = " + dc);
			Debugger.watch("dv = " + dv);
		}
		/*
		 * update derivative of compliance
		 */
		getFilterredComplianceSensitivity(e).setValue(dc);
		/*
		 * update derivative of volume
		 */
		getFilterredVolumeSensitivity(e).setValue(dv);
		/*
		 * save to current variable
		 */
		getTopOptVariable(e).save();
	}

	@Override
	public void FinalizeSolutionStep(Element e) {
		mm.FinalizeSolutionStep(e);
	}

	@Override
	public void CalculateOutputData(Element e) {
		mm.CalculateOutputData(e);
	}

	/**
	 * 
	 * @return the current density change of the material model
	 */
	public double getDensityChange() {
		return this.densityChange;
	}

	@Override
	public IMatrix getConstitutiveMatrix(Element e, int k) {
		IMatrix C = (IMatrix) mm.getConstitutiveMatrix(e, k).clone();
		double xe = getPhysicalDensity(e).getValue();
		double Ee = computeModulus(xe);
		BLAM.scale(Ee / E0, C);
		return C;
	}

	/**
	 * 
	 * @param e
	 * @return topology optimization variable w.r.t element e
	 */
	private TopOptVariable getTopOptVariable(Element e) {
		return elementVariables.get(e);
	}

	/**
	 * get the density of the element with checking the material model ownership
	 * 
	 * @param e
	 * @return
	 */
	public double getElementDensity(Element e) {
		if (this.isPartOf(e)) {
			return getPhysicalDensity(e).getValue();
		} else {
			return -1.0;
		}
	}

	/**
	 * 
	 * @param e
	 * @return physical density of element e
	 */
	protected MutableDouble getPhysicalDensity(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.PHYSICAL_DENSITY);
	}

	/**
	 * 
	 * @param e
	 * @return density of element e
	 */
	protected MutableDouble getDensity(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.DENSITY);
	}

	/**
	 * 
	 * @param e
	 * @return history of density of element e
	 */
	protected double getDensityOld(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.DENSITY, -1);
	}

	/**
	 * 
	 * @param e
	 * @return derivative of compliance w.r.t density
	 */
	protected MutableDouble getFilterredComplianceSensitivity(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.DCDXE);
	}

	/**
	 * 
	 * @param e
	 * @return history of derivative of compliance w.r.t density
	 */
	protected double getComplianceSensitivity(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.DCDXE, -1);
	}

	/**
	 * 
	 * @param e
	 * @return derivative of volume w.r.t density
	 */
	protected MutableDouble getFilterredVolumeSensitivity(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.DVDXE);
	}

	/**
	 * 
	 * @param e
	 * @return history of derivative of volume w.r.t density
	 */
	protected double getVolumeSensitivity(Element e) {
		return getTopOptVariable(e).getScalarVariable(VariableName.DVDXE, -1);
	}

	/**
	 * compute Ee(xe)
	 * 
	 * @param xe
	 * @return
	 */
	private double computeModulus(double xe) {
		return Emin + Math.pow(xe, penal) * (E0 - Emin);
	}

	/**
	 * compute dEe(xe)/dxe
	 * 
	 * @param xe
	 * @return
	 */
	private double computeDModulus(double xe) {
		return penal * Math.pow(xe, penal - 1) * (E0 - Emin);
	}

	// @Override
	// public double computeEnergyAt(IVector e) {
	// return mm.computeEnergyAt(e);
	// }
	//
	// @Override
	// public double computeEnergyAt(IVector e, IVector s) {
	// return mm.computeEnergyAt(e, s);
	// }

	/************************************************
	 * observer overriding: global update for density variables
	 ************************************************/
	@Override
	public void fireStarted() {
		/*
		 * initialize filter
		 */
		filter.initialize();
	}

	@Override
	public void fireTimeStepStarted() {
	}

	@Override
	public void fireIterationStarted() {
	}

	@Override
	public void fireIterationFinished() {

	}

	@Override
	public void fireTimeStepFinished() {
		double l1 = 0.0;
		double l2 = 1e9;
		double move = 0.2;

		/*
		 * filter sensitivities
		 */
		for (Element e : elementVariables.keySet()) {
			if (e.getEchoLevel().contain(EchoLevelType.DEBUG)) {
				Debugger.watch("non-filter dc = " + getComplianceSensitivity(e));
				Debugger.watch("non-filter dv = " + getVolumeSensitivity(e));
			}
			double dc = filter.filterComplianceSensitivity(e, this);
			double dv = filter.filterVolumeSensitivity(e, this);
			getFilterredComplianceSensitivity(e).setValue(dc);
			getFilterredVolumeSensitivity(e).setValue(dv);
			if (e.getEchoLevel().contain(EchoLevelType.DEBUG)) {
				Debugger.watch("filterred dc = " + dc);
				Debugger.watch("filterred dv = " + dv);
			}
		}

		/*
		 * compute new densities by means of bisection algorithm
		 */
		while ((l2 - l1) / (l2 + l1) > MathUtilities.tol3) {
			double lmid = 0.5 * (l1 + l2);
			double actualVolume = 0.0;
			double volume = 0.0;
			for (Element e : elementVariables.keySet()) {
				double dc = getFilterredComplianceSensitivity(e).getValue();
				double dv = getFilterredVolumeSensitivity(e).getValue();
				double xe = getDensityOld(e);

				double xnew = Math.max(
						0.0,
						Math.max(
								xe - move,
								Math.min(1.0,
										Math.min(xe + move, xe * Math.sqrt(-(dc / dv) / lmid)))));
				if (e.getEchoLevel().contain(EchoLevelType.DEBUG)) {
					Debugger.watch("xnew = " + xnew);
				}

				MutableDouble xeNew = getDensity(e);
				xeNew.setValue(xnew);
			}

			/*
			 * filter densitites
			 */
			for (Element e : elementVariables.keySet()) {
				double xnew = filter.filterDensity(e, this);
				getPhysicalDensity(e).setValue(xnew);
				if (e.getEchoLevel().contain(EchoLevelType.DEBUG)) {
					Debugger.watch("filterred xnew = " + xnew);
				}
			}

			/*
			 * check volume and bisect
			 */
			for (Element e : elementVariables.keySet()) {
				double xnew = getPhysicalDensity(e).getValue();
				double ve = e.getMeshPart().computeVolume();
				volume += ve;
				actualVolume += xnew * ve;
			}

			if (actualVolume > volfrac * volume) {
				l1 = lmid;
			} else {
				l2 = lmid;
			}

			// Debugger.watch("l1 = " + l1);
			// Debugger.watch("l2 = " + l2);
		}

		/*
		 * compute change
		 */
		double change = 0.0;
		for (Element e : elementVariables.keySet()) {
			double xnew = getDensity(e).getValue();
			double xe = getDensityOld(e);
			double dx = Math.abs(xnew - xe);
			if (dx > change) {
				change = dx;
			}
			/*
			 * save the variable
			 */
			getTopOptVariable(e).save();
			if (e.getEchoLevel().contain(EchoLevelType.OUTPUT)) {
				xe = getPhysicalDensity(e).getValue();
				Debugger.watch("physical xe = " + xe);
			}
		}
		this.densityChange = change;
		Debugger.watch("density change = " + change);
	}

	@Override
	public void fireFinished() {
	}

}
